from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import AIMessage,HumanMessage,SystemMessage

from langgraph.graph import StateGraph,END
from typing import Annotated,List,TypedDict
import operator

api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

llm = ChatOpenAI(model="gpt-4o",api_key=api_key,base_url=api_base)

class State(TypedDict):
    messages:Annotated[List[str],operator.add]

builder = StateGraph(State)

def chat_with_model(state):
    print(state)
    messages = state["messages"]
    response=llm.invoke(messages)
    return {"messages":[response]}

def convert_messages(state):
    print(state)
    # "您是一位数据提取专家，负责从文本中检索关键信息。请为所提供的文本提取相关信息，并以 JSON 格式输出。概述所提取的关键数据点。"
    EXTRACTION_PROMPT = """
        You are a data extraction specialist tasked with retrieving key information from a text.
        Extract such information for the provided text and output it in JSON format. Outline the key data points extracted.
        """
    messages = [
        SystemMessage(content=EXTRACTION_PROMPT),
        HumanMessage(content=state["messages"][-1].content)
    ]

    response=llm.invoke(messages)
    return {"messages":[response]}

builder.add_node("chat_with_model",chat_with_model)
builder.add_node("convert_messages",convert_messages)

builder.set_entry_point("chat_with_model")
builder.add_edge("chat_with_model","convert_messages")
builder.add_edge("convert_messages",END)

graph=builder.compile()
response = graph.invoke({"messages":[HumanMessage(content="请介绍一下你自己")]})

print(response)